package com.michael.tree.rb;


import static com.michael.tree.rb.RBNode.BLACK;
import static com.michael.tree.rb.RBNode.RED;

/**
 * 红黑树的特性：
 * 1. 节点的颜色只能是红色或者是黑色
 * 2. 根节点的颜色是黑色
 * 3. 叶子节点的颜色只能是黑色
 * 4. 如果一个节点的颜色是红色，则其两个子节点一定都是黑色
 * 5. 对每个节点，从该节点到其所有后代叶节点的所有路径上，都包含数量相同的黑色节点。
 * <p>
 * 进一步解读：
 * 特性1，3基本上不用考虑。
 * 至于根节点的颜色是黑色的，只需要在修复红黑树末尾修改根节点的颜色为黑色即可。所以，该特性2也是不用处理
 * 只有特性4，5是在插入或者删除时有很大可能性被破坏，所以，我们更多情况下，是如何修复该特性。
 */
public class RBTree<T extends Comparable> {
    private RBNode root;


    public void add(T elem) {
        RBNode node = root;
        if (node == null) {
            //如果是空节点，则将其设置为根节点
            root = new RBNode(elem, null);
            return;
        }
        int cmp;
        RBNode parent;
        //合适要插入的父节点
        do {
            parent = node;
            cmp = elem.compareTo(node.key);
            if (cmp < 0)
                node = node.left;
            else if (cmp > 0)
                node = node.right;
            else
                return;
        } while (node != null);
        RBNode newNode = new RBNode(elem, parent);
        if (cmp < 0)
            parent.left = newNode;
        else
            parent.right = newNode;
        fixAfterIntert(newNode);
    }

    public int getHeight() {
        return maxHeight(root);
    }

    public void setRoot(RBNode root) {
        this.root = root;
    }

    public RBNode getRoot() {
        return root;
    }

    private int maxHeight(RBNode node) {
        if (node == null) {
            return 0;
        } else {
            return Math.max(maxHeight(node.left), maxHeight(node.right)) + 1;
        }
    }

    public void remove(T elem) {
        //从树中找到该节点
        RBNode current = getNode(elem);
        if (current == null) {
            return;
        }
        //将当前节点从树中剔除出去
        //如果当前节点有两个子节点
        //则从树中找出比当前节点要大的列表中最小的那个节点，进行替换。
        if (current.left != null && current.right != null) {
            RBNode node = getMinOfGreaterNodesWithNode(current);
            current.key = node.key;
            //将current指向将要替换掉的那个节点
            current = node;
        }
        //运行到这里，当前节点不可能还有两个节点。
        RBNode replace = current.left != null ? current.left : current.right;
        if (replace != null) {
            //如果当前节点还有子节点，则将当前节点剔除出去，将子节点霸占其位置
            if (current.parent == null) {
                root = replace;
            } else if (current == current.parent.left) {
                current.parent.left = replace;
            } else {
                current.parent.right = replace;
            }
            replace.parent = current.parent;
            current.parent = current.right = current.left = null;
            //如果当前节点是黑色的，则就破坏了特性5，需要修复
            if (current.color == BLACK)
                fixAfterDelete(replace);
        } else if (current.parent == null) {
            root = null;
        } else {//当前节点没有子节点
            if (current.color == BLACK) {
                fixAfterDelete(current);
            }
            //将当前节点剔除出去
            RBNode parent = current.parent;
            if (current == parent.left) {
                parent.left = null;
            } else {
                parent.right = null;
            }
            current.parent = null;
        }
    }

    private RBNode getMinOfGreaterNodesWithNode(RBNode node) {
        RBNode min = node.right;
        while (min.left != null) {
            min = min.left;
        }
        return min;
    }

    private RBNode getNode(T elem) {
        RBNode node = root;
        while (node != null) {
            int cmp = node.key.compareTo(elem);
            if (cmp == 0) {
                return node;
            } else if (cmp < 0) {
                node = node.right;
            } else {
                node = node.left;
            }

        }
        return null;
    }

    /**
     * 删除修复，主要是修复特性5：从某一结点到后代节点的黑色子节点的数量是一样的。
     * 当前节点current是没有子节点的
     */
    private void fixAfterDelete(RBNode current) {
        /**
         * while循环是为了将额外的节点沿树上移
         * 额外的节点：基于一种假设的情景下，给当前节点赋予多一个黑色，以至于性质5不会被破坏。
         * 当前节点就是那个额外的节点，
         */
        while (current != root && colorOf(current) == BLACK) {
            if (current == leftOf(parentOf(current))) {
                //兄弟节点是大儿子
                RBNode brother = rightOf(parentOf(current));
                if (colorOf(brother) == RED) {//情况1
                    /**
                     * 当兄弟节点为红色时，其必有两个子节点，都是为黑色。
                     * 因为当前节点是黑色，且还有额外的黑色。其黑色的数量至少是为2的。
                     * 所以，其必有两个子节点。且树的高度会高一点，需要左旋，将树的高度进行降低
                     * 进行左旋，是为了确保兄弟节点的颜色为黑色。才能将额外的节点向上移
                     */
                    setColor(brother, BLACK);
                    setColor(parentOf(current), RED);
                    rotateLeft(parentOf(current));
                    brother = rightOf(parentOf(current));
                }
                //运行到这里，兄弟节点的颜色是黑色
                if (colorOf(leftOf(brother)) == BLACK && colorOf(rightOf(brother)) == BLACK) {//情况2
                    /**
                     * 如果兄弟节点的两个子节点都是黑色，加上当前节点也是黑色(双重黑色)。
                     * 则需要将将兄弟节点设置为红色，然后将额外的节点向上移。
                     */
                    setColor(brother, RED);
                    current = parentOf(current);
                } else {
                    if (colorOf(rightOf(brother)) == BLACK) {//情况3-兄弟节点的大儿子的颜色是黑色
                        /**
                         * 需要将兄弟节点的小儿子的颜色设置为黑色,
                         * 兄弟节点设置为红色
                         * 右旋
                         */
                        setColor(leftOf(brother), BLACK);
                        setColor(brother, RED);
                        rotateRight(brother);
                        brother = rightOf(parentOf(current));
                    }
                    //情况4
                    setColor(brother, colorOf(parentOf(current)));
                    setColor(parentOf(current), BLACK);
                    setColor(rightOf(brother), BLACK);
                    rotateLeft(parentOf(current));
                    current = root;
                }
            } else {
                RBNode brother = leftOf(parentOf(current));
                if (colorOf(brother) == RED) {
                    setColor(brother,BLACK);
                    setColor(parentOf(current),RED);
                    rotateRight(parentOf(current));
                    brother = leftOf(parentOf(current));
                }
                if(colorOf(leftOf(brother))==BLACK && colorOf(rightOf(brother))==BLACK){
                    setColor(brother,RED);
                    current = parentOf(current);
                }else{
                    if(colorOf(leftOf(brother))==BLACK){
                        setColor(brother,RED);
                        setColor(rightOf(brother),BLACK);
                        rotateRight(brother);
                        brother = leftOf(parentOf(current));
                    }
                    setColor(brother,colorOf(parentOf(current)));
                    setColor(parentOf(current),BLACK);
                    setColor(rightOf(brother),BLACK);
                    rotateRight(parentOf(current));
                    current = root;
                }
            }
        }
        setColor(current, RED);
    }

    /**
     * 插入修复
     * <p>
     * 右子节点->大儿子
     * 比父节点还要的叫大舅父，即是当前节点的父亲是祖父的小儿子，所以叔叔节点既是大舅父
     */
    private void fixAfterIntert(RBNode current) {
        current.color = RED;
        //这个时候，从根节点到任意的后代节点的黑色节点的数量是一样的。即特性5不会变，所以就需要判断特性4是否有被破坏
        /**
         * while循环是将红色节点向上移动
         */
        while (current != null && current != root && colorOf(parentOf(current)) == RED) {
            //获取叔叔节点
            /**
             * 为什么获取叔叔节点，个人理解是。
             */
            if (parentOf(current) == leftOf(parentOf(parentOf(current)))) {
                //叔叔节点
                RBNode uncle = rightOf(parentOf(parentOf(current)));
                if (colorOf(uncle) == RED) {//情况1
                    //如果叔叔节点的颜色与父节点的都是红色，
                    //则将叔叔节点与父节点的颜色都置为黑色，这样子父节点的黑色数量与叔叔节点的数量一直，
                    // 同时也不违反特性4
                    setColor(parentOf(current), BLACK);
                    setColor(uncle, BLACK);
                    //我们继续将红色节点,向祖父节点移动。
                    setColor(parentOf(parentOf(current)), RED);
                } else { //叔叔节点为黑色
                    /**
                     * 在这个时候，暂停一下，假设刚刚插入的节点的叔叔节点黑色，而父节点的颜色是红色，这种情况下会出现么？
                     * 答案是不可能出现。仔细想一下
                     * 所以，只有当循环修复关系过程中，才会出现。
                     */
                    if (current == rightOf(parentOf(current))) {//情况2
                        //为什么当叔叔节点是大舅父，当前节点是大儿子时，需要进行左旋呢？
                        //结合情况3的动作，是为了将其将当前节点向上移(移到祖父节点位置上)
                        //设想一下，如果不进行左旋，直接进行右旋，依然无法修复特性4：红色节点的子节点必须是黑色。
                        current = parentOf(current);
                        //以当前节点进行左旋(曾经的父节点，曾经的当前节点作为现在当前节点的父节点)
                        rotateLeft(current);
                    }
                    //情况3
                    //将父节点的颜色置为黑色(如果当前节点是父节点的右子节点的情况下，曾经的祖父节点)
                    setColor(parentOf(current), BLACK);
                    //将祖父节点的颜色置为红色(如果当前节点是父节点的右子节点情况下，曾经的叔叔节点)
                    setColor(parentOf(parentOf(current)), RED);
                    //以祖父节点为基准进行右旋，
                    rotateRight(parentOf(parentOf(current)));
                }
            } else {
                //下面与上面的逻辑差不多,参照上面的解释即可。
                RBNode uncle = leftOf(parentOf(parentOf(current)));
                if (colorOf(uncle) == RED) {
                    setColor(parentOf(current), BLACK);
                    setColor(uncle, BLACK);
                    setColor(parentOf(parentOf(current)), RED);
                    current = parentOf(parentOf(current));
                } else {
                    if (current == leftOf(parentOf(current))) {
                        current = parentOf(current);
                        rotateRight(current);
                    }
                    setColor(parentOf(current), BLACK);
                    setColor(parentOf(parentOf(current)), RED);
                    rotateLeft(parentOf(parentOf(current)));
                }
            }
        }
    }

    /**
     * 左旋,
     * 其将当前节点的右子节点作为当前节点的父节点。
     * 当前节点是右子节点的左子节点
     * 右子节点的左子节点是当前的右子节点
     */
    private void rotateLeft(RBNode node) {
        if (node != null) {
            RBNode parent = node.parent;
            RBNode right = node.right;
            RBNode rightLeft = right.left;
            //修改右子节点的左子节点的指向
            if (rightLeft != null) {
                rightLeft.parent = node;
            }
            node.right = rightLeft;
            //修改父节点指向
            if (parent == null) {
                root = right;
            } else if (parent.left == node) {
                parent.left = right;
            } else {
                parent.right = right;
            }
            right.parent = parent;
            //修改当前子节点的指向
            node.parent = right;
            right.left = node;
        }
    }

    /**
     * 右旋
     * 将当前节点的左子节点作为当前父节点
     * 当前节点是左子节点的右子节点
     * 左子节点的右子节点是当前节点的左子节点
     */
    private void rotateRight(RBNode node) {
        if (node != null) {
            RBNode parent = node.parent;
            RBNode left = node.left;
            RBNode leftRight = left.right;
            //修改左子节点的右子节点的指向
            if (leftRight != null) {
                leftRight.parent = left;
            }
            node.left = leftRight;
            //修改父节点的指向
            if (parent == null) {
                root = left;
            } else if (parent.left == node) {
                parent.left = left;
            } else {
                parent.right = left;
            }
            left.parent = parent;
            //修改当前节点的指向
            node.parent = left;
            left.right = node;
        }
    }

    /***
     *  工具方法,之所以需要这些方法，是为了避免空指针异常
     */
    private static <T extends Comparable> boolean colorOf(RBNode<T> p) {
        return (p == null ? BLACK : p.color);
    }

    private static <T extends Comparable> void setColor(RBNode<T> p, boolean color) {
        if (p != null)
            p.color = color;
    }

    private static <T extends Comparable> RBNode<T> parentOf(RBNode<T> p) {
        return (p == null ? null : p.parent);
    }

    private static <T extends Comparable> RBNode<T> leftOf(RBNode<T> p) {
        return (p == null) ? null : p.left;
    }

    private static <T extends Comparable> RBNode<T> rightOf(RBNode<T> p) {
        return (p == null) ? null : p.right;
    }

}
